import numpy as np

np.random.seed(42)
def generate_random_matrix(rows,cols):
    """生成指定大小的随机矩阵，元素值范围为-1,1"""
    return np.random.uniform(-1,1,(rows,cols))

def quantize_to_int8(matrix):
    scale = np.max(np.abs(matrix))/ 127
    quantized_matrix = np.round(matrix/scale).astype(np.int8)
    return quantized_matrix,scale
# new_matrix = generate_random_matrix(2,2)
# print(new_matrix)
def dequantize_from_int8(quantized_matrix,scale):
    """将量化的矩阵反量化"""
    return quantized_matrix.astype(np.float32)*scale

def quantize_to_int4(matrix):
    """将矩阵量化为int4表示"""
    scale = np.max(np.abs(matrix))/ 7
    quantized_matrix = np.round(matrix / scale)
    quantized_matrix = np.clip(quantized_matrix,-8,7).astype(np.int8)
    return quantized_matrix,scale

def dequantize_from_int4(quantized_matrix,scale):
    return quantized_matrix.astype(np.float32)*scale

def calculate_mean_squared_error(original_matrix,quantized_matrix):
    return np.mean((original_matrix - quantized_matrix)**2)

if __name__ == '__main__':
    np.random.seed(42)
    #生成随机数
    rows,cols = 100,100
    original_matrix = generate_random_matrix(rows,cols)
    # print(np.max(np.abs(original_matrix)))
    quantized_int8,scale_int8 = quantize_to_int8(original_matrix)
    dequantized_int8 = dequantize_from_int8(quantized_int8,scale_int8)
    print(calculate_mean_squared_error(original_matrix,dequantized_int8))

    #int4量化
    quantized_int4,scale_int4 = quantize_to_int4(original_matrix)
    dequantized_int4 = dequantize_from_int4(quantized_int4,scale_int4)
    print(calculate_mean_squared_error(original_matrix,dequantized_int4))